from manim import *

class SieveOfEratosthenes(Scene):
    def construct(self):
        # 标题
        title = Text("埃拉托色尼筛法", font_size=48)
        self.play(Write(title))
        self.wait(1)
        self.play(FadeOut(title))

        # 创建数字网格
        numbers = VGroup(*[Integer(i) for i in range(2, 10)])
        numbers.arrange_in_grid(rows=10, cols=10, buff=0.5)
        self.play(Write(numbers))
        self.wait(1)

        # 筛选质数
        primes = []
        non_primes = set()
        for number in numbers:
            n = number.get_value()
            if n not in non_primes:
                # 高亮当前质数
                self.play(number.animate.set_color(YELLOW))
                primes.append(n)
                
                # 显示当前质数
                prime_text = Text(f"当前质数: {n}", font_size=36).to_edge(UP)
                self.play(Write(prime_text))
                self.wait(1)
                
                # 筛除当前质数的倍数
                for multiple in range(n*2, 10, n):
                    non_primes.add(multiple)
                    # 找到倍数对应的对象
                    multiple_obj = next((num for num in numbers if num.get_value() == multiple), None)
                    if multiple_obj:
                        self.play(multiple_obj.animate.set_color(GRAY))
                        self.wait(0.1)
                
                # 移除当前质数的文本
                self.play(FadeOut(prime_text))

        # 显示找到的质数
        prime_text = Text("100以内的质数", font_size=36).to_edge(UP)
        self.play(Write(prime_text))
        self.wait(1)

        # 高亮所有质数
        highlight_primes = [number.animate.set_color(YELLOW) for number in numbers if number.get_value() in primes]
        self.play(AnimationGroup(*highlight_primes, lag_ratio=0.1))
        self.wait(1)

        # 结束
        end_text = Text("筛选完成！", font_size=48).to_edge(DOWN)
        self.play(Write(end_text))
        self.wait(2)